import React, { useState } from 'react';
import axios from 'axios';
import './StyleTransfer.css';

const StyleTransfer = () => {
  const [contentImage, setContentImage] = useState(null);
  const [styleImage, setStyleImage] = useState(null);
  const [stylizedImage, setStylizedImage] = useState(null);
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState(null);
  const [previewContent, setPreviewContent] = useState(null);
  const [previewStyle, setPreviewStyle] = useState(null);

  const handleContentImageChange = (event) => {
    const file = event.target.files[0];
    if (file) {
      setContentImage(file);
      setPreviewContent(URL.createObjectURL(file));
      setError(null);
    }
  };

  const handleStyleImageChange = (event) => {
    const file = event.target.files[0];
    if (file) {
      setStyleImage(file);
      setPreviewStyle(URL.createObjectURL(file));
      setError(null);
    }
  };

  const handleSubmit = async (event) => {
    event.preventDefault();
    if (!contentImage || !styleImage) {
      setError('请选择内容图像和风格图像');
      return;
    }

    setLoading(true);
    setError(null);
    const formData = new FormData();
    formData.append('content', contentImage);
    formData.append('style', styleImage);

    try {
      const response = await axios.post('http://localhost:5000/style-transfer', formData);
      const { image_url } = response.data;
      setStylizedImage(`http://localhost:5000${image_url}`);
    } catch (error) {
      setError('处理图像时发生错误，请重试');
      console.error('Error processing style transfer:', error);
    } finally {
      setLoading(false);
    }
  };

  const handleReset = () => {
    setContentImage(null);
    setStyleImage(null);
    setStylizedImage(null);
    setPreviewContent(null);
    setPreviewStyle(null);
    setError(null);
  };

  const ImagePreview = ({ src, alt, onDelete }) => (
    <div className="image-preview">
      <img src={src} alt={alt} />
      {onDelete && (
        <button className="delete-button" onClick={onDelete}>
          ×
        </button>
      )}
    </div>
  );

  return (
    <div className="style-transfer-container">
      <div className="style-transfer-content">
        <h1 className="title">图像风格迁移</h1>
        <p className="subtitle">
          上传一张内容图像和一张风格图像，我们将为您生成融合两种风格的图像
        </p>

        {error && (
          <div className="error-message">
            {error}
          </div>
        )}

        <form onSubmit={handleSubmit}>
          <div className="image-grid">
            <div className="image-column">
              <h2>内容图像</h2>
              <label className="upload-button">
                <span>选择内容图像</span>
                <input
                  type="file"
                  accept="image/*"
                  onChange={handleContentImageChange}
                />
              </label>
              {previewContent && (
                <ImagePreview 
                  src={previewContent} 
                  alt="内容图像预览"
                  onDelete={() => {
                    setContentImage(null);
                    setPreviewContent(null);
                  }}
                />
              )}
            </div>

            <div className="image-column">
              <h2>风格图像</h2>
              <label className="upload-button">
                <span>选择风格图像</span>
                <input
                  type="file"
                  accept="image/*"
                  onChange={handleStyleImageChange}
                />
              </label>
              {previewStyle && (
                <ImagePreview 
                  src={previewStyle} 
                  alt="风格图像预览"
                  onDelete={() => {
                    setStyleImage(null);
                    setPreviewStyle(null);
                  }}
                />
              )}
            </div>
          </div>

          <div className="button-group">
            <button
              type="submit"
              className="submit-button"
              disabled={loading || !contentImage || !styleImage}
            >
              {loading ? '处理中...' : '开始风格迁移'}
            </button>
            <button
              type="button"
              className="reset-button"
              onClick={handleReset}
            >
              重置
            </button>
          </div>
        </form>

        {stylizedImage && (
          <div className="result-section">
            <h2>风格迁移结果</h2>
            <div className="result-image-container">
              <img 
                src={stylizedImage} 
                alt="风格迁移结果" 
                className="result-image"
              />
            </div>
          </div>
        )}
      </div>
    </div>
  );
};

export default StyleTransfer; 